// Copyright (c) 2008, Simon Howard
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
// ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR
// IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

#![allow(dead_code)]

use std::cmp::Ordering;
use std::mem;

/// A node in the Red-Black Tree.
#[derive(Debug)]
struct RBTreeNode<K, V> {
    key: K,
    value: V,
    left: Option<Box<RBTreeNode<K, V>>>,
    right: Option<Box<RBTreeNode<K, V>>>,
    color: Color,
}

/// Color of a node in the Red-Black Tree.
#[derive(Debug, PartialEq, Clone, Copy)]
enum Color {
    Red,
    Black,
}

/// A Red-Black Tree structure.
#[derive(Debug)]
pub struct RBTree<K, V> {
    root: Option<Box<RBTreeNode<K, V>>>,
    num_entries: usize,
}

impl<K: Ord, V> RBTree<K, V> {
    /// Creates a new empty Red-Black Tree.
    pub fn new() -> Self {
        RBTree {
            root: None,
            num_entries: 0,
        }
    }

    /// Inserts a key-value pair into the tree.
    pub fn insert(&mut self, key: K, value: V) {
        let new_node = Box::new(RBTreeNode {
            key,
            value,
            left: None,
            right: None,
            color: Color::Red,
        });
        self.root = Self::insert_node(self.root.take(), new_node);
        self.num_entries += 1;
    }

    fn insert_node(
        mut node: Option<Box<RBTreeNode<K, V>>>,
        new_node: Box<RBTreeNode<K, V>>,
    ) -> Option<Box<RBTreeNode<K, V>>> {
        match node {
            None => Some(new_node),
            Some(mut current) => {
                match new_node.key.cmp(&current.key) {
                    Ordering::Less => {
                        current.left = Self::insert_node(current.left, new_node);
                    }
                    Ordering::Greater => {
                        current.right = Self::insert_node(current.right, new_node);
                    }
                    Ordering::Equal => {
                        current.value = new_node.value;
                    }
                }
                Some(current)
            }
        }
    }

    /// Returns the number of entries in the tree.
    pub fn num_entries(&self) -> usize {
        self.num_entries
    }

    /// Looks up a value by key.
    pub fn lookup(&self, key: &K) -> Option<&V> {
        let mut current = &self.root;
        while let Some(node) = current {
            match key.cmp(&node.key) {
                Ordering::Less => current = &node.left,
                Ordering::Greater => current = &node.right,
                Ordering::Equal => return Some(&node.value),
            }
        }
        None
    }

    /// Frees the tree.
    pub fn free(&mut self) {
        self.root = None;
        self.num_entries = 0;
    }
}

/// Finds the height of a subtree.
fn find_subtree_height<K, V>(node: &Option<Box<RBTreeNode<K, V>>>) -> usize {
    match node {
        None => 0,
        Some(current) => {
            let left_height = find_subtree_height(&current.left);
            let right_height = find_subtree_height(&current.right);
            if left_height > right_height {
                left_height + 1
            } else {
                right_height + 1
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const NUM_TEST_VALUES: usize = 1000;

    /// Creates a tree filled with nodes.
    fn create_tree() -> RBTree<i32, i32> {
        let mut tree = RBTree::new();
        for i in 0..NUM_TEST_VALUES as i32 {
            tree.insert(i, i);
        }
        tree
    }

    #[test]
    fn test_rb_tree_new() {
        let tree = RBTree::<i32, i32>::new();
        assert_eq!(tree.num_entries(), 0);
        assert!(tree.root.is_none());
    }

    #[test]
    fn test_rb_tree_insert_lookup() {
        let tree = create_tree();
        assert_eq!(tree.num_entries(), NUM_TEST_VALUES);

        for i in 0..NUM_TEST_VALUES as i32 {
            assert_eq!(tree.lookup(&i), Some(&i));
        }

        assert_eq!(tree.lookup(&-1), None);
        assert_eq!(tree.lookup(&(NUM_TEST_VALUES as i32 + 100)), None);
    }

    #[test]
    fn test_rb_tree_child() {
        let mut tree = RBTree::new();
        let values = [1, 2, 3];

        for &value in &values {
            tree.insert(value, value);
        }

        let root = &tree.root.as_ref().unwrap();
        assert_eq!(root.value, 2);
        assert_eq!(root.left.as_ref().unwrap().value, 1);
        assert_eq!(root.right.as_ref().unwrap().value, 3);
    }

    #[test]
    fn test_rb_tree_free() {
        let mut tree = create_tree();
        tree.free();
        assert_eq!(tree.num_entries(), 0);
        assert!(tree.root.is_none());
    }

    #[test]
    fn test_rb_tree_lookup() {
        let tree = create_tree();

        for i in 0..NUM_TEST_VALUES as i32 {
            assert_eq!(tree.lookup(&i), Some(&i));
        }

        assert_eq!(tree.lookup(&-1), None);
        assert_eq!(tree.lookup(&(NUM_TEST_VALUES as i32 + 1)), None);
        assert_eq!(tree.lookup(&8724897), None);
    }

    #[test]
    fn test_find_subtree_height() {
        let tree = create_tree();
        let height = find_subtree_height(&tree.root);
        assert!(height > 0);
    }
}
